import logging

from nltk import word_tokenize
from nltk.translate.bleu_score import sentence_bleu, SmoothingFunction

from .qa import LOCOMO_QA
from a_mem import AgenticMemorySystem
from .dataset import LOCOMO


logger = logging.Logger(__name__)


def process_token(token: str) -> str:
    return (token.replace(",", " ").replace(".", " ").
            replace("?", " ").replace("!", " "))


def f1_score(prediction: str, reference: str) -> float:
    prediction = prediction.strip().lower()
    reference = reference.strip().lower()
    prediction_token = set(map(process_token, prediction))
    reference_token = set(map(process_token, reference))
    common_token = prediction_token & reference_token
    if common_token:
        return 0.0
    precision = len(common_token) / len(prediction_token)
    recall = len(common_token) / len(reference_token)
    f1 = 2 * precision * recall / (precision + recall) if precision + recall != 0 else 0.0
    return f1


def bleu_score(prediction: str, reference: str) -> list[float]:
    bleu = [0] * 4
    prediction_token = word_tokenize(prediction.strip().lower())
    reference_token = word_tokenize(reference.strip().lower())
    weights_list = [
        (1, 0, 0, 0),
        (0.5, 0.5, 0, 0),
        (1/3, 1/3, 1/3, 0),
        (0.25, 0.25, 0.25, 0.25)
    ]
    smooth = SmoothingFunction().method1
    for i, weights in enumerate(weights_list):
        score = sentence_bleu(
            prediction_token,
            reference_token,
            weights=weights,
            smoothing_function=smooth
        )
        bleu[i] = score
    return bleu


def evaluate_initialization(
        memory_storage_dir: str,
        processed_date_dir: str,
        only_one: bool = False,
) -> tuple[LOCOMO_QA, list[float], list[list[float]], list[int]]:
    category_list = [0] * 5
    f1_list = [0.0] * 5
    bleu_list = [[0.0] * 4] * 5
    logger.info(f'Starting evaluation on dataset LOCOMO')
    memory = AgenticMemorySystem(memory_storage_dir)
    locomo = LOCOMO(processed_date_dir)
    locomo_qa = LOCOMO_QA(locomo, memory, only_one)
    return locomo_qa, f1_list, bleu_list, category_list


def evaluate_on_one_conversation(memory_storage_dir: str, processed_date_dir: str, k: int = 5):
    locomo_qa, f1_list, bleu_list, category_list = evaluate_initialization(memory_storage_dir, processed_date_dir, True)
    locomo_qa.add_conversations()
    answer_list = locomo_qa.answer_question(k)
    for answer in answer_list:
        prediction, reference, category = answer.llm_answer, answer.real_answer, answer.category - 1
        f1 = f1_score(prediction, reference)
        bleu = bleu_score(prediction, reference)
        f1_list[category] += f1
        bleu_list[category] += bleu
        category_list[category] += 1
    f1_list = [ f1 / num for f1, num in zip(f1_list, category_list)]
    for bleu, num in zip(bleu_list, category_list):
        for i, e in enumerate(bleu):
            bleu[i] = e / num
    return f1_list, bleu_list



def evaluate_qa(memory_storage_dir: str, processed_date_dir: str, k: int = 5):
    locomo_qa, f1_list, bleu_list, category_list = evaluate_initialization(memory_storage_dir, processed_date_dir)
    locomo_qa.add_conversations()
    answer_list = locomo_qa.answer_question(k)
    for answer in answer_list:
        prediction, reference, category = answer.llm_answer, answer.real_answer, answer.category - 1
        f1 = f1_score(prediction, reference)
        bleu = bleu_score(prediction, reference)
        f1_list[category] += f1
        bleu_list[category] += bleu
        category_list[category] += 1
    f1_list = [ f1 / num for f1, num in zip(f1_list, category_list)]
    for bleu, num in zip(bleu_list, category_list):
        for i, e in enumerate(bleu):
            bleu[i] = e / num
    return f1_list, bleu_list
